import csv
import random
import argparse
from itertools import permutations

def generate_nodes(output_file, num_nodes):
    """生成随机节点数据"""
    with open(output_file, 'w', newline='') as f:
        writer = csv.DictWriter(f, fieldnames=["id", "lon", "lat"])
        writer.writeheader()
        for node_id in range(1, num_nodes + 1):
            lon = round(random.uniform(-180, 180), 6)
            lat = round(random.uniform(-90, 90), 6)
            writer.writerow({"id": str(node_id), "lon": str(lon), "lat": str(lat)})

def generate_edges(output_file, num_edges, num_nodes):
    """生成随机唯一边数据（无自环、无重复）"""
    all_possible_edges = set(
        (str(from_id), str(to_id))
        for from_id, to_id in permutations(range(1, num_nodes + 1), 2)
        if from_id != to_id
    )
    
    if num_edges > len(all_possible_edges):
        raise ValueError(f"最多只能生成 {len(all_possible_edges)} 条唯一边（当前请求 {num_edges} 条）")
    
    sampled_edges = random.sample(sorted(all_possible_edges), num_edges)
    
    with open(output_file, 'w', newline='') as f:
        writer = csv.DictWriter(f, fieldnames=["from", "to"])
        writer.writeheader()
        for from_node, to_node in sampled_edges:
            writer.writerow({"from": from_node, "to": to_node})

def main():
    parser = argparse.ArgumentParser(description="生成随机节点和边的 CSV 数据")
    parser.add_argument("--nodes", type=int, default=100, help="生成的节点数（默认：100）")
    parser.add_argument("--edges", type=int, default=150, help="生成的边数（默认：150）")
    parser.add_argument("--nodes_file", default="nodes.csv", help="节点输出文件（默认：nodes.csv）")
    parser.add_argument("--edges_file", default="edges.csv", help="边输出文件（默认：edges.csv）")
    args = parser.parse_args()

    if args.nodes < 1:
        raise ValueError("节点数至少为 1")
    if args.edges < 1:
        raise ValueError("边数至少为 1")

    generate_nodes(args.nodes_file, args.nodes)
    generate_edges(args.edges_file, args.edges, args.nodes)
    print(f"生成完成！节点数：{args.nodes}，边数：{args.edges}")

if __name__ == "__main__":
    main()